import mysql.connector
from sqlalchemy.ext.compiler import compiles
from mysql.connector import Error
from contextlib import contextmanager
from app.config import Config
from app.utils.logger import logger
from sqlalchemy import create_engine, Table, MetaData
from sqlalchemy.sql.expression import Insert
from sqlalchemy.dialects.mysql import insert
from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine, text
import time 
from urllib.parse import quote_plus
class MySQLHelper:
    def __init__(self):
        self.host = Config.DATABASE_HOST
        self.port = Config.DATABASE_PORT
        self.user = Config.DATABASE_USER
        self.password = quote_plus(Config.DATABASE_PASSWORD) 
        self.db = Config.DATABASE_DATABASE
        self.charset = "utf8mb4"
        self.engine = None
        self.Session = None
        self._create_engine()

    def _create_engine(self):
        # 创建数据库连接引擎，使用连接池
        # 检查端口是否为空，如果为空则使用默认端口3306
        port = self.port if self.port else "3306"
        # 使用与配置文件一致的连接器，默认使用mysqlconnector
        connection_string = f"mysql+mysqlconnector://{self.user}:{self.password}@{self.host}:{port}/{self.db}?charset={self.charset}"
        self.engine = create_engine(connection_string, pool_size=10, max_overflow=1490, pool_recycle=3600)
        self.Session = sessionmaker(bind=self.engine)

    def get_session(self):
        # 获取数据库会话
        return self.Session()

    def mysql_replace_into(self, table, conn, keys, data_iter):
        # 定义一个编译器扩展，用于将 INSERT INTO 替换为 REPLACE INTO
        @compiles(Insert)
        def replace_string(insert, compiler, **kw):
            s = compiler.visit_insert(insert, **kw)
            s = s.replace("INSERT INTO", "REPLACE INTO")
            return s

        # 准备数据
        data = [dict(zip(keys, row)) for row in data_iter]
        
        # 执行插入操作
        conn.execute(table.table.insert(), data)

    def _df_res_to_db(self, pd, table_name):
        pd.to_sql(
            name=table_name,
            con=self.engine.connect(),
            index=False,
            chunksize=2000,
            if_exists='append', method=self.mysql_replace_into)
            
    def execute_select(self, sql):
        time_start = time.time()
        with self.engine.connect() as conn:
            # 使用 text() 函数将 SQL 语句转换为可执行对象
            sql_text = text(sql)
            result = conn.execute(sql_text)
            data = []
            for row in result:
                # 使用 row._mapping 将 RowProxy 转换为字典
                data.append(dict(row._mapping))
            time_end = time.time()
            # print(f"执行时间：{time_end - time_start}")
            return data

    def execute_value(self, sql):
        with self.engine.connect() as conn:
            result = conn.execute(sql)
            return result.fetchone()

    def execute_edit(self, sql):
        with self.engine.connect() as conn:
            result = conn.execute(sql)
            conn.commit()
            return result.rowcount
            
    def execute_charts(self, sql):
        """
        :return 图表数据
        """
        with self.engine.connect() as conn:
            result = conn.execute(sql)
            target = {}
            for j in range(len(result.keys())):
                data = [row[j] for row in result]
                target[result.keys()[j]] = data
            return target

    def execute_datatable(self, sql):
        with self.engine.connect() as conn:
            result = conn.execute(sql)
            datas = []
            for row in result:
                target = list(row)
                datas.append(target)
            data = {"data": datas}
            return data

    def execute_cloud(self, sql):
        with self.engine.connect() as conn:
            result = conn.execute(sql)
            data = {row[0]: row[1] for row in result}
            return data
            
    def get_exhibition_by_id(self, exhibition_id):
        """
        根据ID获取展会信息
        """
        if not exhibition_id:
            return None
            
        sql = f"""
        SELECT *
        FROM exhibitions 
        WHERE id = {exhibition_id}
        LIMIT 1
        """
        
        result = self.execute_select(sql)
        if result and len(result) > 0:
            return result[0]
        return None
        
    def update_exhibition(self, exhibition_id, update_data):
        """
        更新展会信息
        
        Args:
            exhibition_id: 展会ID
            update_data: 需要更新的字段和值的字典
        """
        if not exhibition_id or not update_data:
            return False
            
        # 构建SET部分的SQL
        set_clauses = []
        for field, value in update_data.items():
            if isinstance(value, str):
                # 字符串需要转义并加引号
                value = value.replace("'", "''")
                set_clauses.append(f"{field} = '{value}'")
            elif value is None:
                set_clauses.append(f"{field} = NULL")
            else:
                set_clauses.append(f"{field} = {value}")
                
        if not set_clauses:
            return False
            
        set_sql = ", ".join(set_clauses)
        
        sql = f"""
        UPDATE exhibitions 
        SET {set_sql}
        WHERE id = {exhibition_id}
        """
        
        # 使用事务执行更新
        with self.engine.connect() as conn:
            with conn.begin():
                result = conn.execute(text(sql))
                return result.rowcount > 0
        
    def update_ai_status(self, exhibition_id, status):
        """
        更新展会的AI补全状态
        
        Args:
            exhibition_id: 展会ID
            status: 状态描述
        """
        if not exhibition_id:
            return False
            
        # 处理状态文本，防止SQL注入
        if isinstance(status, str):
            status = status.replace("'", "''")
            
        sql = f"""
        UPDATE exhibitions 
        SET ai_checked_status = 
            CASE 
                WHEN '{status}' = 'AI补全中' THEN 4
                WHEN '{status}' = 'AI补全成功' THEN 5
                WHEN '{status}' LIKE 'AI补全失败%' THEN 2
                ELSE 3
            END,
            ai_check_description = '{status}'
        WHERE id = {exhibition_id}
        """
        
        # 使用事务执行更新
        with self.engine.connect() as conn:
            with conn.begin():
                result = conn.execute(text(sql))
                return result.rowcount > 0

class Database:
    """数据库连接管理类"""
    @staticmethod
    def get_connection():
        try:
            # 检查端口是否为空，如果为空则使用默认端口3306
            port = Config.DATABASE_PORT if Config.DATABASE_PORT else "3306"
            
            conn = mysql.connector.connect(
                host=Config.DATABASE_HOST,
                port=port,
                user=Config.DATABASE_USER,
                password=Config.DATABASE_PASSWORD,
                database=Config.DATABASE_DATABASE,
                autocommit=False  # 禁用自动提交
            )
            return conn
        except Error as e:
            logger.error(f"数据库连接失败: {e}")
            return None

    @contextmanager
    def get_cursor(self, dictionary=True):
        """上下文管理器：自动管理连接和游标"""
        try:
            if not self._conn:
                self._conn = self.get_connection()
            if not self._cursor:
                self._cursor = self._conn.cursor(dictionary=dictionary)
            
            # 重写 cursor.execute 方法来记录 SQL
            original_execute = self._cursor.execute
            def execute_wrapper(query, params=None):
                if params:
                    logger.info(f"执行SQL: {query}, 参数: {params}")
                else:
                    logger.info(f"执行SQL: {query}")
                return original_execute(query, params)
            self._cursor.execute = execute_wrapper
            
            yield self._cursor
            
            # 直接提交，让事务管理由外层控制
            # if self._conn:
            #     self._conn.commit()
                
        except Error as e:
            if self._conn:
                self._conn.rollback()
            logger.error(f"数据库操作失败: {e}")
            raise Exception(f"数据库操作失败: {e}")

    def execute_query(self, query, params=None, dictionary=True):
        """执行查询并返回结果"""
        with self.get_cursor(dictionary=dictionary) as cursor:
            cursor.execute(query, params or ())
            logger.info(f"执行查询: {query}")
            if params:
                logger.info(f"参数: {params}")
            result = cursor.fetchall()
            logger.info(f"查询结果数: {len(result)}")
            return result

    def execute_update(self, query, params=None):
        """执行更新并返回受影响的行数"""
        with self.get_cursor() as cursor:
            cursor.execute(query, params or ())
            logger.info(f"执行更新: {query}")
            if params:
                logger.info(f"参数: {params}")
            affected_rows = cursor.rowcount
            logger.info(f"影响行数: {affected_rows}")
            # 如果不在显式事务中，自动提交
            if not self._in_transaction:
                self._conn.commit()
            return affected_rows

    def __init__(self):
        self._conn = None
        self._cursor = None
        self._in_transaction = False  # 添加事务状态标记

    def begin_transaction(self):
        """开始事务"""
        try:
            if not self._conn:
                self._conn = self.get_connection()
            if self._conn:
                if not self._in_transaction:  # 检查是否已在事务中
                    self._in_transaction = True
                    logger.info("事务开始")
                else:
                    logger.warning("事务已经在进行中，跳过开启操作")
            else:
                raise Exception("无法建立数据库连接")
        except Exception as e:
            logger.error(f"开始事务失败: {str(e)}")
            raise e

    def commit(self):
        """提交事务"""
        try:
            if self._conn and self._in_transaction:
                self._conn.commit()
                self._in_transaction = False
        except Exception as e:
            logger.error(f"提交事务失败: {str(e)}")
            raise e

    def rollback(self):
        """回滚事务"""
        try:
            if self._conn and self._in_transaction:
                self._conn.rollback()
                self._in_transaction = False
        except Exception as e:
            logger.error(f"回滚事务失败: {str(e)}")
            raise e

    def close(self):
        """关闭数据库连接"""
        try:
            if self._cursor:
                self._cursor.close()
            if self._conn:
                if self._in_transaction:
                    self._conn.rollback()  # 如果还在事务中，先回滚
                self._conn.close()
        finally:
            self._cursor = None
            self._conn = None
            self._in_transaction = False